from contextlib import contextmanager
from typing import Callable, TYPE_CHECKING, Union

if TYPE_CHECKING:
    from pyarrow.parquet import ParquetFile


@contextmanager
def parquet_file_source_context(source: Union[Callable[[], 'ParquetFile'], 'ParquetFile']):
    if isinstance(source, Callable):
        with source() as parquet_file:
            yield parquet_file
    else:
        parquet_file = source
        yield parquet_file
